# %%
import cv2
import os
input_path = "../../assets/video/truck/"  #需要处理的视频所在的目录
output_path="../../assets/image/truck/" #处理后的图片存放目录


file_names = [f for f in os.listdir(input_path) if f.endswith(".mp4")]
# %%
def extract_frames(video_path, output_dir, format='jpg', frame_interval=1,rescnt=0):
    """
    从视频中提取帧并保存为图像
    :param video_path: 视频文件路径
    :param output_dir: 输出目录
    :param format: 图像格式 ('jpg' 或 'png')
    :param frame_interval: 每隔多少帧保存一 (默认为1,即每帧都保存）
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("Error: 无法打开视频文件！")
        return
    
    frame_count = 0
    saved_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break  # 视频结束
     
        if frame_count % frame_interval == 0:   # 按间隔保存帧
            if format.lower() == 'png':
                output_path = os.path.join(output_dir, f"frame_{saved_count+rescnt:05d}.png")
                cv2.imwrite(output_path, frame, [cv2.IMWRITE_PNG_COMPRESSION, 3])  # PNG压缩级别（0-9）
            elif format.lower() == 'jpg':
                output_path = os.path.join(output_dir, f"frame_{saved_count+rescnt:05d}.jpg")
                cv2.imwrite(output_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 95])  # JPG质量（0-100）
            else:
                raise ValueError("不支持的格式，请选择 'jpg' 或 'png'")
            saved_count += 1
        
        frame_count += 1
    
    cap.release()
    return saved_count

total=0
for file in file_names:
    # 使用示例
   saved=extract_frames(
        video_path=input_path+file,      # 输入视频路径
        output_dir=output_path,  # 输出目录
        format="jpg",               # 保存为JPG格式
        frame_interval=20,           # 每10帧保存1张
        rescnt=total
    )
   assert saved is not None
   total+=saved
print(f"提取完成！保存 {total} 张图像到 {output_path}")


